#include "sgPardiso.h"
#include <iostream>
#include "mkl.h"
#include "mkl_pardiso.h"
#include "mkl_sparse_handle.h"
//#include "mkl_spblas.h"

/*
 * The pardiso routine calculates the solution of a set of sparse linear equations A*X = B
 * with single or multiple right-hand sides, using a parallel LU, LDL, or LLT factorization,
 * where A is an n-by-n matrix, and X and B are n-by-nrhs vectors or matrices.
 *
 * see https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2024-2/pardiso.html
 *
 * Relatated Codes from oneMKL docs, which are under the following directory
 * /opt/Intel/oneAPI/mkl/2024.2/share/doc/mkl/examples/examples_core_c/c/sparse_directsolvers/source
 *
 * - pardiso_sym.c
 * Example program to show the use of the "PARDISO" routine
 * on symmetric linear systems
 * which is under /opt/Intel/oneAPI/mkl/2024.2/share/doc/mkl/examples/examples_core_c/c/sparse_directsolvers/source/pardiso_sym.c
 *
 * - pardiso_unsym_csc.c
 * Intel(R) oneAPI Math Kernel Library (oneMKL) PARDISO C example program to show the
 * use of the "PARDISO" routine on unsymmetric linear systems in CSC format.
 */

//
// Several key notes are shown as the followings,
//
// 1. The matrix storage in Eigen
// The SparseMatrix and SparseVector classes take three template arguments:
// the scalar type (e.g., double) the storage order (ColMajor or RowMajor,
// the default is ColMajor) the inner index type (default is int).
// see https://eigen.tuxfamily.org/dox/group__TutorialSparse.html
//
// 2. Dump Eigen::SparseMatrix into PARDISO
// iparm[36]: Format for matrix storage.
// Use CSR format (see Three Array Variation of CSR Format) for matrix storage.
// see https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2024-2/pardiso-iparm-parameter.html
// The 3-array variation of the CSR format has a restriction: all non-zero elements
// are stored continuously, that is the set of non-zero elements in the row J goes
// just after the set of non-zero elements in the row J-1.
// see https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2024-2/sparse-blas-csr-matrix-storage-format.html#CSR3
//
// iparm[11] = 2: Solve a transposed system ATX = B based on the factorization of
// the matrix A.
//
// On output, some iparm values report information such as the numbers of non-zero
// elements in the factors.
// see https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2024-2/pardiso.html
//
// 3. Compatibility with MinGW-w64
// MinGW is not a compiler platform officially supported for MKL. MKL is only compatable
// with Visual studio compatiable library, include some MSVC run-time library.
//
// see https://community.intel.com/t5/Intel-oneAPI-Math-Kernel-Library/MKL-with-MinGW64-compiler/m-p/802929#M3221%3Fwapkw=mingw64
//

// Define the format to printf MKL_INT values
#if !defined(MKL_ILP64)
#define IFORMAT "%i"
#else
#define IFORMAT "%lli"
#endif

// _DEBUG Defined as 1 when the /LDd, /MDd, or /MTd compiler option is set.
// Otherwise, undefined.
// see https://learn.microsoft.com/en-us/cpp/preprocessor/predefined-macros?view=msvc-170#standard-predefined-identifier

#ifdef SG_VERBOSE
#define BUFLEN 16
// The mkl_progress function is intended to track progress of a lengthy
// computation and/or interrupt the computation.  A non-zero return value
// may be supplied by the redefined function to break the computation.
// While a user-supplied mkl_progress function usually redefines the
// default mkl_progress function automatically, some configurations require
// calling the mkl_set_progress function to replace the default mkl_progress
// function.
// see
// https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2024-2/mkl-progress.html#GUID-E982C8B0-E96D-414B-B263-BABF455963F7

int mkl_progress (int* thread_process, int* step, char* stage, int lstage)
{
    char buf[BUFLEN];
    if (lstage >= BUFLEN)
        lstage = BUFLEN - 1;
    strncpy (buf, stage, lstage);
    buf[lstage] = '\0';
    printf ("In thread %i, at stage %s, steps passed %i\n", *thread_process, buf, *step);
    return 0;
}

#endif

using namespace SG::Algebra;

namespace SG::Algebra
{
    template <typename TDataType>
    struct TPardisoPrivate
    {
        /* Matrix data in CSC format. */
        /* Number of equations in the sparse linear systems of equations A*X = B.*/
        // MKL_INT n = 0;
        // MKL_INT* ia = nullptr;
        // MKL_INT* ja = nullptr;
        // double*              a = nullptr;
        std::vector<TDataType> a;
        std::vector<MKL_INT> ia;
        std::vector<MKL_INT> ja;

        MKL_INT mtype = 11; /* Real unsymmetric matrix */
        // // Descriptor of main sparse matrix properties
        // struct matrix_descr descrA;
        // // Structure with sparse matrix stored in CSR format
        // sparse_matrix_t csrA;
        MKL_INT nrhs = 1; /* Number of right hand sides. */
        /* Internal solver memory pointer pt, */
        /* 32-bit: int pt[64]; 64-bit: long int pt[64] */
        /* or void *pt[64] should be OK on both architectures */
        void* pt[64];

        /* Pardiso control parameters. */
        MKL_INT iparm[64];
        MKL_INT maxfct, mnum, error, msglvl;
        /* Auxiliary variables. */
        TDataType ddum; /* Double dummy */
        MKL_INT idum; /* Integer dummy. */
    };
}  // namespace Algebra

//=============================================================================
// TPardiso

template <typename TDataType>
TPardiso<TDataType>::TPardiso ()
{
    d = new TPardisoPrivate<TDataType> ();
    // test_pardiso_unsym_csc ();
    init ();
}

template <typename TDataType>
TPardiso<TDataType>::~TPardiso ()
{
    clearFactors ();
    if (d)
        delete d;
}

template <typename TDataType>
void TPardiso<TDataType>::init ()
{
    m_factorizationIsOk = false;

    if (typeid (double) == typeid (TDataType))
    {
        d->mtype = 11; /* Real unsymmetric matrix */
    }
    else if (typeid (Complex_t) == typeid (TDataType))
    {
        d->mtype = 13; /* Complex unsymmetric matrix */
    }

    d->nrhs = 1;       /* Number of right hand sides. */

    /* -------------------------------------------------------------------- */
    /* .. Setup Pardiso control parameters. */
    /* -------------------------------------------------------------------- */
    for (int i = 0; i < 64; i++)
    {
        d->iparm[i] = 0;
    }
    d->iparm[0]  = 1;  /* No solver default */
    d->iparm[1]  = 2;  /* Fill-in reordering from METIS */
    d->iparm[3]  = 0;  /* No iterative-direct algorithm */
    d->iparm[4]  = 0;  /* No user fill-in reducing permutation */
    d->iparm[5]  = 0;  /* Write solution into x */
    d->iparm[6]  = 0;  /* Not in use */
    d->iparm[7]  = 2;  /* Max numbers of iterative refinement steps */
    d->iparm[8]  = 0;  /* Not in use */
    d->iparm[9]  = 13; /* Perturb the pivot elements with 1E-13 */
    d->iparm[10] = 1;  /* Use nonsymmetric permutation and scaling MPS */
    d->iparm[11] = 0;  /* Conjugate/transpose solve */
    d->iparm[12] = 1;  /* Maximum weighted matching algorithm is switched-on (default for non-symmetric) */
    d->iparm[13] = 0;  /* Output: Number of perturbed pivots */
    d->iparm[14] = 0;  /* Not in use */
    d->iparm[15] = 0;  /* Not in use */
    d->iparm[16] = 0;  /* Not in use */
    d->iparm[17] = -1; /* Output: Number of nonzeros in the factor LU */
    d->iparm[18] = -1; /* Output: Mflops for LU factorization */
    d->iparm[19] = 0;  /* Output: Numbers of CG Iterations */
    d->maxfct    = 1;  /* Maximum number of numerical factorizations. */
    d->mnum      = 1;  /* Which factorization to use. */
    //d->msglvl    = 0;  /* pardiso generates no output  */
    d->msglvl    = 0;  /* pardiso generates no output  */
    d->error     = 0;  /* Initialize error flag */

    /* -------------------------------------------------------------------- */
    /* .. Initialize the internal solver memory pointer. This is only */
    /* necessary for the FIRST call of the PARDISO solver. */
    /* -------------------------------------------------------------------- */
    for (int i = 0; i < 64; i++)
    {
        d->pt[i] = 0;
    }

#ifdef SG_VERBOSE
    d->msglvl    = 1;  /* Print statistical information  */
    d->iparm[26] = 1;  // Matrix checker.
#endif

    // d->iparm[11] = 2;  /*Solve a transposed system ATX = B based on the factorization of the matrix A*/
    // d->iparm[33] = 1;  // Optimal number of OpenMP threads for conditional numerical reproducibility (CNR) mode.
    d->iparm[34] = 1; /* Zero-based indexing: columns and rows indexing in arrays ia, ja, and perm starts from 0 (C-style indexing).*/
}

template <typename TDataType>
void TPardiso<TDataType>::clearFactors ()
{
    if (!m_factorizationIsOk)  // decompsed?
    {
        return;
    }
    MKL_INT n = d->ia.size () - 1;

    // mkl_sparse_destroy (d->csrA);

    /* -------------------------------------------------------------------- */
    /* .. Termination and release of memory. */
    /* -------------------------------------------------------------------- */
    MKL_INT phase = -1; /* Release internal memory. */
    PARDISO (d->pt,
             &d->maxfct,
             &d->mnum,
             &d->mtype,
             &phase,
             &n,
             &d->ddum,
             d->ia.data (),
             d->ja.data (),
             &d->idum,
             &d->nrhs,
             d->iparm,
             &d->msglvl,
             &d->ddum,
             &d->ddum,
             &d->error);

    // flag that LU factorization has been not performed.
    m_factorizationIsOk = false;
}

template <typename TDataType>
void TPardiso<TDataType>::compute (const SparseMatrixType& matrix)
{
    // the solver can be re-used to another matrix equations
    // Thus, clear the old LU factorization.
    clearFactors ();

    init ();

    SparseMatrixType& mat = const_cast<SparseMatrixType&> (matrix);

    MKL_INT m   = mat.rows ();
    MKL_INT n   = mat.cols ();
    MKL_INT nnz = mat.nonZeros ();
    // dump index/values into pardiso
    // d->a = mat.valuePtr ();
    d->a.resize (nnz);
    for (auto i = 0; i < nnz; ++i)
    {
        d->a[i] = *(mat.valuePtr () + i);
    }

    d->ia.resize (m + 1);
    for (auto i = 0; i < m + 1; ++i)
    {
        d->ia[i] = mat.outerIndexPtr ()[i];
    }

    d->ja.resize (nnz);
    for (auto i = 0; i < nnz; ++i)
    {
        d->ja[i] = mat.innerIndexPtr ()[i];
    }

    // checkMatrix ();

    /* -------------------------------------------------------------------- */
    /* .. Reordering and Symbolic Factorization. This step also allocates */
    /* all memory that is necessary for the factorization. */
    /* -------------------------------------------------------------------- */
    MKL_INT phase = 11;
    PARDISO (d->pt,
             &d->maxfct,
             &d->mnum,
             &d->mtype,
             &phase,
             &n,
             &d->a[0],
             &d->ia[0],
             &d->ja[0],
             &d->idum,
             &d->nrhs,
             d->iparm,
             &d->msglvl,
             &d->ddum,
             &d->ddum,
             &d->error);
    if (d->error != 0)
    {
        printf ("\nERROR during symbolic factorization: \n" IFORMAT, d->error);
        exit (1);
    }
    // printf ("\nReordering completed ... ");
    // printf ("\nNumber of nonzeros in factors = " IFORMAT, d->iparm[17]);
    // printf ("\nNumber of factorization MFLOPS = " IFORMAT, d->iparm[18]);

    /* -------------------------------------------------------------------- */
    /* .. Numerical factorization. */
    /* -------------------------------------------------------------------- */
    phase = 22;
    PARDISO (d->pt,
             &d->maxfct,
             &d->mnum,
             &d->mtype,
             &phase,
             &n,
             &d->a[0],
             &d->ia[0],
             &d->ja[0],
             &d->idum,
             &d->nrhs,
             d->iparm,
             &d->msglvl,
             &d->ddum,
             &d->ddum,
             &d->error);
    if (d->error != 0)
    {
        printf ("\nERROR during numerical factorization: \n" IFORMAT, d->error);
        exit (2);
    }
    // printf ("Factorization completed ... \n");

    // Now, the LU factorization has been finished.
    m_factorizationIsOk = true;
}

template <typename TDataType>
void TPardiso<TDataType>::solve (const VectorType& B, VectorType& X) const
{
    MKL_INT m = B.rows ();
    MKL_INT n = B.rows ();

    // On entry, contains the right-hand side vector/matrix B, which is placed in memory contiguously.
    // The b[+k*nrhs] element must hold the i-th component of k-th right-hand side vector.
    VectorType b = B;

    // If iparm[5]=0 it contains solution vector/matrix X, which is placed contiguously in memory.
    // The x[i + k*n] element must hold the i-th component of the k-th solution vector.
    VectorType x (m);

    /* -------------------------------------------------------------------- */
    /* .. Back substitution and iterative refinement. */
    /* -------------------------------------------------------------------- */
    MKL_INT phase = 33;

    // d->descrA.type = SPARSE_MATRIX_TYPE_GENERAL;
    // d->descrA.mode = SPARSE_FILL_MODE_UPPER;
    // d->descrA.diag = SPARSE_DIAG_NON_UNIT;

    // // rows_start[i] - indexing is the first index of row i in the arrays values and col_indx.
    // // rows_end[i] - indexing - 1 is the last index of row i in the arrays values and col_indx.
    // // For zero-based indexing, array containing the column indices for each non-zero element of the matrix A.
    // // Array containing non-zero elements of the matrix A
    // // see https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2024-2/mkl-sparse-create-csr.html
    // MKL_INT* rows_start = d->ia.data ();
    // MKL_INT* rows_end   = d->ia.data () + 1;
    // MKL_INT* col_indx   = d->ja.data();
    // double*  values     = d->a;

    // mkl_sparse_d_create_csr (&d->csrA, SPARSE_INDEX_BASE_ZERO, n, n, rows_start, rows_end, col_indx, values);

    // d->iparm[5]  = 1;  /* The solver stores the solution on the right-hand side b. */
    d->iparm[11] = 2;  // Transpose solve is used for systems in CSC format

    // printf ("\n\nSolving the system in CSC format...\n");
    PARDISO (d->pt,
             &d->maxfct,
             &d->mnum,
             &d->mtype,
             &phase,
             &n,
             &d->a[0],
             &d->ia[0],
             &d->ja[0],
             &d->idum,
             &d->nrhs,
             d->iparm,
             &d->msglvl,
             b.data (),
             x.data (),
             &d->error);
    if (d->error != 0)
    {
        printf ("\nERROR during solution: \n" IFORMAT, d->error);
        exit (3);
    }

    // fetch the solution
    X.resize (m);
    for (auto i = 0; i < m; ++i)
    {
        X[i] = x[i];
    }
}

template <typename TDataType>
void TPardiso<TDataType>::solve(const SparseMatrixType& B, SparseMatrixType& X) const 
{
    MKL_INT m = B.rows();
    MKL_INT n = B.cols(); 
    d->nrhs = 1;

    // Check if the coefficient matrix A is a square matrix
    MKL_INT A_n = d->ia.size() - 1;
    if (A_n != m) {
        printf("\nERROR: A.rows() (%d) != B.rows() (%d)\n", A_n, m);
        exit(2);
    }

    // Pre allocation result matrix storage structure
    SparseMatrixType tempX;
    tempX.resize(m, n);
    std::vector<Eigen::Triplet<TDataType>> triplets;  
    triplets.reserve(B.nonZeros());  // Pre allocate memory based on non-zero input matrix values

    // 单列求解缓冲区
    std::vector<TDataType> rhs(m);   // Store vectors on the right side of each column
    std::vector<TDataType> sol(m);   // Store single column solution results

    for (int i = 0; i < n; ++i)
    {
        // Initialize the right-hand vector (convert sparse columns into dense vectors)
        rhs.assign(m, 0);  // Clear buffer zone
        //create the rhs
        Eigen::Map<VectorType>(rhs.data(), m) = B.col(i); // Extract column data from matrix B
        // Call PARDISO to solve a single column problem
        MKL_INT phase = 33;
        PARDISO(d->pt, 
                &d->maxfct, 
                &d->mnum, 
                &d->mtype, 
                &phase,
                &m, 
                d->a.data(),
                d->ia.data(), 
                d->ja.data(),
                &d->idum, 
                &d->nrhs, 
                d->iparm, 
                &d->msglvl,
                rhs.data(), 
                sol.data(), 
                &d->error);

        if (d->error != 0) {
            printf ("\nERROR during numerical factorization: \n" IFORMAT, d->error);
            exit (3);
        }        
        // store the solution into X
        // Collect non-zero solution elements 
        for (int j = 0; j < m; ++j) { 
            if (std::abs(sol[j]) > 1e-12) {   
                triplets.emplace_back(j, i, sol[j]);
            }
        }
    }
    
    // Build the final sparse matrix
    tempX.setFromTriplets(triplets.begin(), triplets.end());
    tempX.makeCompressed();
    X = tempX.transpose();  
    X.makeCompressed();    
}


template <typename TDataType>
Info_t TPardiso<TDataType>::info () const
{
    return ( Info_t ) d->error;
}

//=============================================================================
// Explicit instantiations to hide the underlying implementation.
namespace SG
{
    namespace Algebra
    {
        template class ALGEBRA_EXPORT TPardiso<Real_t>;
        template class ALGEBRA_EXPORT TPardiso<Complex_t>;
    }  // namespace Algebra
}  // namespace SG

//=============================================================================
// Explicit specilizations

//------------------------------------------------------------------
// Real_t

template <>
const char* ALGEBRA_EXPORT TPardiso<Real_t>::getClassName ()
{
    return "TPardiso<Real_t>";
}

//------------------------------------------------------------------
// Complex_t
template <>
const char* ALGEBRA_EXPORT TPardiso<Complex_t>::getClassName ()
{
    return "TPardiso<Complex_t>";
}